"""
中间件

注意点：在中间件中手动抛出的异常，需要在中间件中手动捕获，无法在异常处理中捕获
"""

import random
import traceback
import uuid
from json import JSONDecodeError

from fastapi import FastAPI
from fastapi.routing import APIRoute
from starlette.middleware.base import (
    BaseHTTPMiddleware,
    RequestResponseEndpoint,
)
from starlette.middleware.cors import CORSMiddleware
from starlette.requests import Request
from starlette.responses import Response

import utils.enum_utils as enum
from config import settings
from models.other import (
    RequestClient,
    RequestInfo,
    RequestRoute,
    ResponseInfo,
)
from models.po import (
    RequestLogPO,
    UserPO,
)
from my_typing import *
from utils.data_utils.special import ObjectIdStr
from utils.exception_utils import AuthenticationError
from utils.mock_utils import DATETIME


class RequestUtils:
    """ 请求相关的工具类 """

    RECORD_CONTENT_TYPES = (
        enum.ContentType.JSON,
    )
    """ 需要记录响应数据的 Content-Type """

    @staticmethod
    def gen_request_id() -> str:
        """
        生成请求 ID

        生成规则（总长度 28 位）：
            - 24 位的 ObjectId
            - 随机的 3 位数
        """

        return "_".join(
            (
                ObjectIdStr(),
                str(random.randint(0, 999)).zfill(3),
            )
        )

    @staticmethod
    def gen_trace_id() -> str:
        """
        生成链路 ID

        生成规则（总长度 40~41 位）：
            - 项目编号
            - 当前环境
            - 24 位的 ObjectId
            - UUID 前 8 位

        Examples:
            >>> RequestUtils.gen_trace_id()
            "ex.dev.67dccecefb7902999354bbdc.af1b86b1"
        """

        return ".".join(
            (
                settings.CONST.PROJECT_NAME,
                settings.CONST.ENV,
                ObjectIdStr(),
                uuid.uuid4().hex[:8],
            )
        ).lower()

    @staticmethod
    async def async_data_generator(async_generator: AsyncGenerator) -> tuple[AsyncGenerator, list[bytes]]:
        """ 取出异步生成器的内容，并返回异步生成器、内容列表 """

        data_list: list = [
            chunk
            async for chunk in async_generator
        ]

        async def async_data_generator() -> AsyncGenerator:
            for i in data_list:
                yield i

        return async_data_generator(), data_list

    @staticmethod
    async def get_method_with_default(
        request: Request,
        method_name: str,
        default: T,
    ) -> T:
        """ 获取请求方法，如果发生异常则返回默认值 """

        from components.logger import log

        result = default
        default_type = type(default)

        try:
            method = getattr(request, method_name, default_type)

            if method is not default_type:
                result = default_type(await method())

        except (JSONDecodeError, UnicodeDecodeError):  # noqa
            ...

        except Exception as _:  # noqa: ignore[broad exception]
            log.error(f"[系统][获取请求方法异常]: {traceback.format_exc()}")

        return result

    @staticmethod
    def get_route_info(request: Request) -> RequestRoute:
        """ 获取路由信息 """

        route_info: dict = {
            "method": request.method,
            "path": request.url.path,
        }
        route: APIRoute | None = None

        for r in request.app.routes:
            if (
                isinstance(r, APIRoute)
                and r.path_regex.match(request.url.path)
                and request.method in r.methods
            ):
                route = r
                break

        if route:
            route_info = {
                **route.__dict__,
                "dependencies": [
                    _.dependency.__name__
                    for _ in route.dependencies
                ],
                "deprecated": bool(route.deprecated),
                "path_regex": str(route.path_regex),
                "request_model": (
                    str(route.body_field.type_)
                    if route.body_field
                    else None
                ),
                "response_model": str(route.response_model),
                **route_info,
            }

        return RequestRoute(**route_info)

    @classmethod
    def parse_form_file(cls, form: dict) -> dict:
        """ 解析 form 文件，转换成 dict """

        return {
            k: str(v)
            for k, v in form.items()
        }

    @classmethod
    async def get_request_info(cls, request: Request) -> RequestInfo:
        """ 获取请求信息 """

        body = await cls.get_method_with_default(request, "json", {})
        form = await cls.get_method_with_default(request, "form", {})
        form = cls.parse_form_file(form)

        return RequestInfo(
            client=RequestClient(
                host=request.client.host,
                port=request.client.port,
            ),
            headers=dict(request.headers),
            body=body,
            params=dict(request.query_params),
            form=form,
        )

    @classmethod
    async def get_response_info(cls, response: Response) -> ResponseInfo:
        """ 获取响应信息 """

        response_data: ListStr = []

        if (
            (body_iterator := getattr(response, "body_iterator", None))
            and response.headers.get("content-type") in cls.RECORD_CONTENT_TYPES
        ):
            generator, data_list = await cls.async_data_generator(body_iterator)
            setattr(response, "body_iterator", generator)

            for _ in data_list:
                response_data.append(_.decode("UTF-8"))

        return ResponseInfo(
            headers=dict(response.headers),
            body=response_data,
            status_code=response.status_code,
            media_type=response.media_type,
        )


class GlobalParamMiddleware(BaseHTTPMiddleware):
    """
    全局参数 | 中间件

        - 请求 ID
        - 链路追踪 ID
    """

    async def dispatch(
        self,
        request: Request,
        call_next: RequestResponseEndpoint,
    ) -> Response:
        from components.global_param import (
            update_request_context,
            update_request_scope,
        )

        # 获取或生成 trace_id
        trace_id: str = request.headers.get("X-Trace-ID") or RequestUtils.gen_trace_id()

        # 生成请求 ID
        request_id: str = RequestUtils.gen_request_id()
        request_route: RequestRoute = RequestUtils.get_route_info(request)

        # 使用 ContextVar 存储请求信息,每次请求创建新的字典
        ctx_data: dict = {
            "request_id": request_id,
            "request_ip": request.client.host,
            "request_time": DATETIME.now(),
            "request_route": request_route,
            "trace_id": trace_id,
        }
        update_request_context(ctx_data)
        update_request_scope(request, ctx_data)

        # 调用下一个中间件
        response = await call_next(request)

        # 添加 request_id、trace_id 到响应头
        response.headers["Request-ID"] = request_id
        response.headers["X-Trace-ID"] = trace_id

        return response


class GlobalExceptionMiddleware(BaseHTTPMiddleware):
    """
    全局异常处理 | 中间件

    注意点：在中间件中手动抛出的异常，需要在中间件中手动捕获，无法在异常处理中捕获
    """

    async def dispatch(self, request, call_next):
        try:
            response = await call_next(request)

        except Exception as _:  # noqa: ignore[broad exception]
            from components.logger import log
            from components.response import ExceptionResponse

            log.error(f"[系统][全局异常]: {traceback.format_exc()}")
            response = ExceptionResponse()

        return response


class AuthenticationMiddleware(BaseHTTPMiddleware):
    """
    鉴权 | 中间件

    获取请求的用户，添加到 request.scope["user"]、request.user、当前线程 中
    """

    async def dispatch(self, request, call_next):
        from components.authentication import get_request_user
        from components.global_param import (
            update_request_context,
            update_request_scope,
        )

        try:
            user_po: UserPO | None = get_request_user(request)

        except AuthenticationError as exc:
            from components.response import ErrorResponse

            return ErrorResponse(exc.errors, code=enum.HTTPStatusCode.HTTP_401_UNAUTHORIZED)

        if user_po:
            # 隐藏密码
            user_po.password = ""

        ctx_data: dict = {"user": user_po}
        update_request_context(ctx_data)
        update_request_scope(request, ctx_data)

        return await call_next(request)


class RequestLogMiddleware(BaseHTTPMiddleware):
    """ 请求日志 | 中间件 """

    async def dispatch(self, request, call_next):
        from biz import request_log_biz

        current_user: UserPO | None = request.scope.get("user")
        request_time: datetime = request.scope.get("request_time") or DATETIME.now()
        route_info: RequestRoute = RequestUtils.get_route_info(request)
        request_info: RequestInfo = await RequestUtils.get_request_info(request)
        response = await call_next(request)
        response_info: ResponseInfo = await RequestUtils.get_response_info(response)
        response_time: datetime = DATETIME.now()
        time_info: dict = {
            "in_time": request_time,
            "out_time": response_time,
            "total_seconds": (response_time - request_time).total_seconds(),
        }
        request_log_po: RequestLogPO = RequestLogPO(
            **{
                "request_id": request.scope["request_id"],
                "request_info": request_info,
                "response_info": response_info,
                "route_info": route_info,
                "time_info": time_info,
                "trace_id": request.scope["trace_id"],
                "user_id": getattr(current_user, "id", None),
            },
        )
        request_log_biz.RequestLogSetter.async_create_request_log(request_log_po)

        return response


def register_middleware(app: FastAPI) -> None:
    """
    注册中间件

    注意点：最后 app.add_middleware 添加的中间件，会最先执行
    """

    # 自定义中间件
    middleware_list = [
        GlobalParamMiddleware,
        GlobalExceptionMiddleware,
        AuthenticationMiddleware,
        RequestLogMiddleware,
    ]

    # 注册自定义中间件
    for _middleware in middleware_list[::-1]:
        app.add_middleware(_middleware)  # noqa: ignore[arg-type]

    # 注册 CORS，最外层
    app.add_middleware(
        CORSMiddleware,  # noqa: ignore[arg-type]
        allow_origins=[
            "http://localhost:5173",
            "http://127.0.0.1:5173",
        ],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )
